import os

from langchain.chat_models import QianfanChatEndpoint
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, PromptTemplate

"""用 PromptTemplate 来生成Prompt"""
# 可以借助PromptTemplate来生成对话的Prompt。下面是一个例子：想要将一种语言翻译成另一种语言。

os.environ["QIANFAN_AK"] = "KBe3HbStsjvnGS9KFrNZFPes"
os.environ["QIANFAN_SK"] = "3yGoXGIdCCy8QLwcLCLdh3UGGlM7PKNn"

chatllm = QianfanChatEndpoint()

# 设置好模板
system_template = '将{input}翻译成{output}'
system_message_prompt_template = SystemMessagePromptTemplate.from_template(system_template)
human_template = '{text}'
human_message_prompt_template = HumanMessagePromptTemplate.from_template(human_template)

# ChatPromptTemplate
chat_prompt_template = ChatPromptTemplate.from_messages([system_message_prompt_template, human_message_prompt_template])
messages = chat_prompt_template.format_prompt(input='中文', output='英语', text='苏州是一座美丽的城市').to_messages()

result = chatllm(messages)
print(result)

# 直接使用PromptTemplate来创建：
prompt_template = PromptTemplate(
    template='将{input}翻译成{output}',
    input_variables=['input', 'output']
)
message_prompt_template = SystemMessagePromptTemplate(prompt=prompt_template)
